from my_rnn import MyRNN
from my_gru import MyGRU
from my_lstm import MyLSTM
import torch
import torch.nn as nn
import string
import glob
import os
import unicodedata
import random
import time
import math
import torch
import matplotlib.pyplot as plt


all_letters = string.ascii_letters + " .,;'"

n_letters = len(all_letters)


def unicode_to_ascii(s):
    """函数的作用是去掉一些语言中的重音标记"""
    return ''.join(c for c in unicodedata.normalize('NFD', s)
                   if unicodedata.category(c) != 'Mn'
                   and c in all_letters)


def read_lines(filename):
    # 打开指定的文件并读取所有的内容, 使用strip()去除掉两侧的空白符, 然后以'\n'为换行符进行切分
    lines = open(filename, encoding='utf-8').read().strip().split('\n')
    return [unicode_to_ascii(line) for line in lines]


# print("n_categories:", n_categories)

# print(category_lines['Italian'][:10])
# exit(0)
def line_to_tensor(line):
    # 首先初始化一个全0的张量, 这个张量的形状是(len(line), 1, n_letters)
    # 代表人名中的每一个字母都用一个(1 * n_letters)张量来表示
    tensor = torch.zeros(len(line), 1, n_letters)
    # 遍历每个人名中的每个字符, 并搜索其对应的索引, 将该索引位置置1
    for li, letter in enumerate(line):
        tensor[li][0][all_letters.find(letter)] = 1

    return tensor


def category_from_output(all_categories, output):
    """
    从输出结果中得到指定的类别, topk函数
    :param output:
    :return:
    """
    topn, topi = output.topk(1)
    category_i = topi[0].item()
    # 从前面已经构造好的all_categories中得到对应语言的类别，返回类别和索引
    return all_categories[category_i], category_i


def random_train_data(all_categories, category_lines):
    """
    该函数用于随机取得训练数据
    :return:
    """
    category = random.choice(all_categories)
    line = random.choice(category_lines[category])

    category_tensor = torch.tensor([all_categories.index(category)], dtype=torch.long)
    line_tensor = line_to_tensor(line)
    return category, line, category_tensor, line_tensor


def init_data():
    # 构建一个人名类别与具体人名对应关系的字典
    category_lines = {}
    # 构建所有类别的列表
    all_categories = []
    data_path = "../data/names/"
    # 遍历所有的文件, 使用glob.glob中可以利用正则表达式的便利
    for filename in glob.glob(data_path + "*.txt"):
        # 获取每个文件的文件名, 起始就是得到名字的类别
        category = os.path.splitext(os.path.basename(filename))[0]
        # 逐一将其装入所有类别的列表中
        all_categories.append(category)
        # 然后读取每个文件的内容, 形成名字的列表
        lines = read_lines(filename)
        # 按照对应的类别, 将名字列表写入到category_lines字典中
        category_lines[category] = lines
    return all_categories, category_lines


def time_since(since):
    now = time.time()
    s = now - since
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)


if __name__ == '__main__':
    print(1)
    exit(0)
    since = time.time() - 8 * 60
    period = time_since(since)
    print(period)
    # exit(0)
    s = "Ślusàrski"
    a = unicode_to_ascii(s)
    print(a)
    input1 = line_to_tensor('Bobby')
    print(input1, input1.shape)
    input1 = input1.squeeze(0)  # squeeze如果不带参数的话，会将所有长度为1的维度给抹掉 [1,1,16] -> [16]
    print(input1, input1.shape)
    exit(0)

    # print("n_letters:", n_letters)
    # exit(0)

    # line = "Bai"
    # line_tensor = line_to_tensor(line)
    # print("line_tensor:", line_tensor)

    all_categories, category_lines = init_data()
    input_size = len(all_letters)
    n_hidden = 128
    output_size = len(all_categories)


    hidden = c = torch.zeros(1, 1, n_hidden)
    rnn = MyRNN(input_size, n_hidden, output_size)
    lstm = MyLSTM(input_size, n_hidden, output_size)
    gru = MyGRU(input_size, n_hidden, output_size)

    rnn_output, next_hidden = rnn(input1, hidden)
    print('rnn:', rnn_output)
    print('rnn_shape:', rnn_output.shape)
    print('***********')

    lstm_output, next_hidden1, c = lstm(input1, hidden, c)
    print('lstm:', lstm_output)
    print('lstm_shape:', lstm_output.shape)
    print('***********')

    gru_output, next_hidden2 = gru(input1, hidden)
    print('gru:', gru_output)
    print('gru_shape:', gru_output.shape)

    category, category_i = category_from_output(gru_output)
    print('category: {}, category_i: {}'.format(category, category_i))

    for i in range(10):
        category, line, category_tensor, line_tensor = random_train_data(all_categories, category_lines)
        print('category:{}, line:{}, category_tensor:{}'.format(category, line, category_tensor))
